백준 2143번

오늘 풀어본 문제는 백준의 2143번 문제¹이다. 문제 풀이에 사용한 언어는 C++ 이다.

solved.ac 기준 CLASS

CLASS 5

문제 정보

이 문제의 내용과 조건은 다음과 같다.

문제

한 배열 $A[1]$, $A[2]$, $\cdots$, $A[n]$ 에 대해서, 부 배열은 $A[i]$, $A[i+1]$, $\cdots$, $A[j-1]$, $A[j]$ (단, $1 \le i \le j \le n$)을 말한다. 이러한 부 배열의 합은 $A[i]+\cdots+A[j]$ 를 의미한다. 각 원소가 정수인 두 배열 $A[1], \cdots, A[n]$ 과 $B[1], \cdots, B[m]$ 이 주어졌을 때, $A$ 의 부 배열의 합에 $B$ 의 부 배열의 합을 더해서 $T$ 가 되는 모든 부 배열 쌍의 개수를 구하는 프로그램을 작성하시오.

예를 들어 $A = {1, 3, 1, 2}$, $B = {1, 3, 2}$, $T=5$ 인 경우, 부 배열 쌍의 개수는 다음의 $7$ 가지 경우가 있다.

T(=5) = A[1] + B[1] + B[2]
      = A[1] + A[2] + B[1]
      = A[2] + B[3]
      = A[2] + A[3] + B[1]
      = A[3] + B[1] + B[2]
      = A[3] + A[4] + B[3]
      = A[4] + B[2] 

입력

첫째 줄에 $T$ $(-1,000,000,000 \le T \le 1,000,000,000)$ 가 주어진다. 다음 줄에는 $n$ $(1 \le n \le 1,000)$ 이 주어지고, 그 다음 줄에 $n$ 개의 정수로 $A[1], \cdots, A[n]$ 이 주어진다. 다음 줄에는 $m$ $(1 \le m \le 1,000)$ 이 주어지고, 그 다음 줄에 $m$ 개의 정수로 $B[1], \cdots, B[m]$ 이 주어진다. 각각의 배열 원소는 절댓값이 $1,000,000$ 을 넘지 않는 정수이다.

출력

첫째 줄에 답을 출력한다. 가능한 경우가 한 가지도 없을 경우에는 $0$ 을 출력한다.

풀이과정

1번째 시도

문제를 해결하기 위해 사용한 방법은 다음과 같다.

반복문을 활용하여 가능한 모든 부분 수열의 합을 각각 구한 뒤 각각 저장한다.
한쪽 부분 수열의 합들을 정렬하고 (여기서는 $B$ 의 부분 수열의 합들을 정렬했다.) 나머지 부분 수열들의 합을 순회하며 더했을 때 $T$ 가 되는 값들의 개수를 이분 탐색을 활용하여 알아낸다.

코드는 다음과 같이 작성하였다.

#include <bits/stdc++.h>

using namespace std;

using ll = long long;
using ull = unsigned long long;

using pii = pair<int, int>;
using pll = pair<ll, ll>;

int main(void) {
    ios::sync_with_stdio(false);
    cin.tie(nullptr);
    cout.tie(nullptr);

    ll T, n, m;
    cin >> T;

    cin >> n;
    vector<ll> A(n);

    for (int i=0; i<n; i++) {
        cin >> A[i];
    }

    cin >> m;
    vector<ll> B(m);

    for (int i=0; i<m; i++) {
        cin >> B[i];
    }

    vector<ll> subSeqSumA;
    vector<ll> subSeqSumB;

    for (int i=0; i<n; i++) {
        ll sum = A[i];
        subSeqSumA.emplace_back(sum);

        for (int j=i+1; j<n; j++) {
            sum += A[j];
            subSeqSumA.emplace_back(sum);
        }
    }

    for (int i=0; i<m; i++) {
        ll sum = B[i];
        subSeqSumB.emplace_back(sum);

        for (int j=i+1; j<m; j++) {
            sum += B[j];
            subSeqSumB.emplace_back(sum);
        }
    }

    sort(B.begin(), B.end());

    ll result = 0;
    for (auto& sumA : subSeqSumA) {
        ll target = T - sumA;

        result += upper_bound(subSeqSumB.begin(), subSeqSumB.end(), target)
                - lower_bound(subSeqSumB.begin(), subSeqSumB.end(), target);;
    }

    cout << result;

    return 0;
}

실행 결과 ‘틀렸습니다’ 가 떴다.

2번째 시도

틀렸던 원인은 subSeqSumB 대신 B 를 정렬한 것이었다. subSeqSumB 가 정렬되지 않아 이분 탐색이 제대로 이루어지지 않았던 것이다.

코드는 다음과 같이 수정하였다.

#include <bits/stdc++.h>

using namespace std;

using ll = long long;
using ull = unsigned long long;

using pii = pair<int, int>;
using pll = pair<ll, ll>;

int main(void) {
    ios::sync_with_stdio(false);
    cin.tie(nullptr);
    cout.tie(nullptr);

    ll T, n, m;
    cin >> T;

    cin >> n;
    vector<ll> A(n);

    for (int i=0; i<n; i++) {
        cin >> A[i];
    }

    cin >> m;
    vector<ll> B(m);

    for (int i=0; i<m; i++) {
        cin >> B[i];
    }

    vector<ll> subSeqSumA;
    vector<ll> subSeqSumB;

    for (int i=0; i<n; i++) {
        ll sum = A[i];
        subSeqSumA.emplace_back(sum);

        for (int j=i+1; j<n; j++) {
            sum += A[j];
            subSeqSumA.emplace_back(sum);
        }
    }

    for (int i=0; i<m; i++) {
        ll sum = B[i];
        subSeqSumB.emplace_back(sum);

        for (int j=i+1; j<m; j++) {
            sum += B[j];
            subSeqSumB.emplace_back(sum);
        }
    }

    sort(subSeqSumB.begin(), subSeqSumB.end());

    ll result = 0;
    for (auto& sumA : subSeqSumA) {
        ll target = T - sumA;

        result += upper_bound(subSeqSumB.begin(), subSeqSumB.end(), target)
                - lower_bound(subSeqSumB.begin(), subSeqSumB.end(), target);;
    }

    cout << result;

    return 0;
}

그러자 모든 테스트 케이스를 통과하고 ‘맞았습니다’가 나오는 것을 확인할 수 있었다.

마무리

이분 탐색 문제가 나올 때면 항상 C++ 을 써서 다행이라는 생각이 든다. lower_bound 와 upper_bound 함수가 미리 구현되어 있기 때문에 그냥 가져다 쓸 수 있기 때문이다. 직접 만들어야 했다면 인덱스 신경쓰느라 매번 귀찮았을 것이다. 물론 한 번만 구현해놓으면 계속 쓸 수 있었겠지만 말이다.

오늘의 PS는 여기까지!

1: https://www.acmicpc.net/problem/2143

← Previous Post Next Post →